# a pytorch based lisv2 code
# author: zelinzang
# email: zangzelin@gmail.com

import numpy as np
import torch
import torch.autograd
import scipy

from sklearn import datasets, manifold
from sklearn.manifold import SpectralEmbedding
from sklearn.metrics.pairwise import euclidean_distances, pairwise_distances
from torch import nn, set_flush_denormal
from torch.autograd import Variable

# from pynndescent import NNDescent
import torch.nn.functional as F

# from torch.nn.modules import loss
# import functools
from multiprocessing import Pool

from torch.nn.modules.conv import Conv2d
import torch
# import torchvision.transforms as transforms

import torch.nn as nn
from torch.nn.modules.batchnorm import _BatchNorm
import GCN
# from mmcv.cnn import constant_init, kaiming_init, normal_init


def savefig(tensor_image, path):
    import matplotlib.pyplot as plt

    plt.imshow(tensor_image.cpu().permute(1, 2, 0).numpy())
    plt.savefig(path)
    plt.close()


class LISV2_CNN(nn.Module):
    def __init__(
        self,
        device,
        args,
        data=None,
        vs=0.001,
        ve=100,
        perplexity=10,
        etapow=2,
        dataset_name="swishroll",
        # NetworkStructure=[-1, 500, 500, 2],
        BitofEmbedding=128,
        CNN_arch="resnet18",  # add new resnet
        path="./",
        n_dim=2,
    ):
        nn.Module.__init__(self)
        self.args = args
        self.device = device
        # self.NetworkStructure = NetworkStructure
        self.BitofEmbedding = BitofEmbedding
        # self.NetworkStructure[0] = data.shape[1]

        # self.data = data.to(self.device)

        self.CNN_arch = CNN_arch
        self._InitNetwork(Cnn=True)

    def _InitNetwork(self, Cnn=False):

        import torchvision.models as models

        resnet = models.__dict__[self.CNN_arch]
        self.encoder = resnet(num_classes=self.BitofEmbedding)
        self.encoder = torch.nn.Sequential(*(list(self.encoder.children())[:-1]))
        # print(self.encoder)
        # input()

    def forward(self, x):

        # x = torch.cat([inputs, ax])
        if self.CNN_arch != "LeNet":
            x = self.encoder(x)
        else:
            for i, layer in enumerate(self.encoder):
                x = layer(x)

        # print(x.shape)
        return x.reshape((x.shape[0], x.shape[1]))

    def test(self, input_data):
        x = input_data.to(self.device)

        if self.CNN_arch != "LeNet":
            x = self.encoder(x)
        else:
            for i, layer in enumerate(self.encoder):
                x = layer(x)

        return [x]

    def GetInput(
        self,
    ):
        return None


class LISV2_MLP(nn.Module):
    def __init__(
        self,
        input_dim,
        networkStructure=[1000, 500, 200, 2],
        NNmethod=''
    ):
        nn.Module.__init__(self)
        # self.BitofEmbedding = BitofEmbedding
        networkStructure[0] = input_dim
        self.NNmethod = NNmethod
        self.networkStructure = networkStructure
        self.InitNetwork(self.networkStructure)
        self.GCL = GCN.GraphConvolution2(250,250)
        # self.GCL = GCN.GraphConvolution2(200,200)

    def InitNetwork(self, NetworkStructure):
        self.encoder = nn.ModuleList()
        for i in range(len(NetworkStructure) - 1):
            self.encoder.append(
                nn.Linear(NetworkStructure[i],
                          NetworkStructure[i + 1]))
            if i != len(NetworkStructure) - 2:
                self.encoder.append(nn.LeakyReLU(0.1))

        self.decoder = nn.ModuleList()
        for i in range(len(NetworkStructure) - 1, 0, -1):
            self.decoder.append(
                nn.Linear(NetworkStructure[i],
                          NetworkStructure[i - 1]))
            if i != 1:
                self.decoder.append(nn.LeakyReLU(0.1))
        # Map output to range (0, 1) for image datasets
        # if('mnist' in self.dataset_name or 'coil' in self.dataset_name):
        #     self.decoder.append(nn.Sigmoid())

    def forward(self, x, adj=None):

        for i, layer in enumerate(self.encoder):
            x = layer(x)
            if i == 3:
                if self.NNmethod=='gcn+NN':
                    # print('ddd')
                    xN = self.GCL(x, adj)
                    x = torch.cat((xN,x), dim=1)
                if self.NNmethod=='gcn':
                    x = self.GCL(x, adj)

        return x

    def test(self, input_data):
        x = input_data.to(self.device)


        for i, layer in enumerate(self.encoder):
            x = layer(x)

        return [x]

    def GetInput(
        self,
    ):
        return None